package com.cyt.base.aspect;

import com.cyt.base.exception.base.BaseException;
import com.cyt.base.util.ErrorCode;
import com.cyt.base.util.ServletUtils;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.context.annotation.Scope;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletRequest;
import java.lang.reflect.Method;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 限流
 *
 * @author tao cong cong
 * @create 2019/3/17 12:14 PM
 */
@Component
@Scope
@Aspect
@Slf4j
public class LimitAspect {

    /**
     * 每秒只发出5个令牌
     * 此处是单进程服务的限流,内部采用令牌捅算法实现
     */
    // private static RateLimiter rateLimiter = RateLimiter.create(1.0);
    private final Map<String, RateLimiter> limitMap = new ConcurrentHashMap<>();

    @Pointcut("@annotation(com.cyt.base.aspect.Limit)")
    public void ServiceAspect() {

    }

    @Around("ServiceAspect()")
    public Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        HttpServletRequest request = ServletUtils.getRequest();
        // url(存在map集合的key)
        String url = request.getRequestURI();
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        if (method.isAnnotationPresent(Limit.class)) {
            Limit limit = method.getAnnotation(Limit.class);
            RateLimiter limiter = null;
            if (!limitMap.containsKey(url)) {
                // 创建令牌桶
                limiter = RateLimiter.create(limit.value());
                limitMap.put(url, limiter);
            }
            limiter = limitMap.get(url);
            // 默认毫秒
            boolean acquire = limiter.tryAcquire(limit.timeOut(), limit.timeOutUnit());
            if (!acquire) {
                // 直接抛出系统繁忙的异常
                log.info("限流，系统繁忙,请稍后再试");
                throw new BaseException(ErrorCode.CODE_1, "系统繁忙,请稍后再试");
            }
        }
        return joinPoint.proceed();
    }

}
